
#!/usr/bin/python
# -*- coding: UTF-8 -*-

import time
import datetime
import argparse
import os
import json
import matplotlib.pyplot as plt

"""
解析日志，可视化性能图片
"""

def str2ms(timestamp):
    """日期格式转换

    Args:
        timestamp: 字符串格式日期（带毫秒）， "%Y-%m-%d %H:%M:%S.%f"
    Return:
        日期的毫秒格式
    """
    fmt = "%Y-%m-%d %H:%M:%S.%f" # %f 是当做微秒解析的
    d1 = datetime.datetime.strptime(timestamp, fmt)
    d2 = time.mktime(d1.timetuple()) # 这个是秒为单位的
    # 转毫秒
    return int(d2 * 1000 + d1.microsecond / 1000)


def ms2str(timestamp):
    """日期格式转换

    Args:
        timestamp: 日期的毫秒格式
    Return: 
        字符串格式日期，"%Y-%m-%d %H:%M:%S"
    """


    # 记录毫秒
    ms = timestamp % 1000
    # 获得秒级时间戳
    fmt = f"%Y-%m-%d %H:%M:%S"
    timeStr = time.strftime(fmt, time.localtime(timestamp / 1000))
    return  f"{timeStr}.{ms:03d}"

def getRange(line, prefix, suffix):
    """截取字符串某段内容

    prefix[这里时结构内容]suffix

    Args:
        line: 字符串

    Return:
        如果存在pattern则返回截取侯的字符串, 否则返回None
    """
    start = line.find(prefix)
    if start == -1:
        print(f"{prefix} 不存在 Fail on: {line}")
        return None
    end = line.find(suffix, start)
    if end == -1:
        print(f"{suffix} 不存在 Fail on: {line}")
        return None
    return line[start+len(prefix):end]

def getStartName(endName):
    """获取匹配的StartName

        XXXEnd -> XXXStart
    """
    if "End" in endName:
        return f"{endName[:endName.find("End")]}Start"
    return None

isDetail = True # 全局Log开关

class TaskStage:
    """
        表示一个任务的某个阶段：
        规则:
            任务开始: Start -> TimeStamp
            某个阶段开始：xxxStart -> TimeStamp
            某个阶段结束：xxxEnd -> TimeStamp
            任务结束：End -> TimeStamp
        
        name: 阶段名称
        timestamp: 该Stage的起始时间戳，统一为毫秒数，int类型
        cost: 当前阶段对应的耗时， int类型， 单位ms， Start阶段是0
        totalCost: 当前阶段相对于任务开始的耗时，int类型， 单位ms
        pid: 当前阶段打点时在哪个pid上执行

    """
    def __init__(self, name, timestamp, pid):
        self.name = name
        self.stamp = timestamp
        self.cost = 0
        self.totalCost = 0
        self.pid = pid

    def isEnd(self):
        return "End" in self.name
    
    def isTaskEnd(self):
        return "End" == self.name
    
    def getStageName(self):
        """获取Stage名字
        """
        if "End" == self.name:
            return "ALL"
        if "End" in self.name:
            return self.name[:self.name.find("End")]
        
    def getStart(self):
        """计算开始时间点的浮点格式: xx.xxxs
        """
        res = float(ms2str(self.stamp - self.cost)[-6:])
        # TODO: bugfix 任务在 [aa:xx.xxx, aa+1:xx.xxx] 要给aa+1的时间戳 + 60
        # 如何感知到分钟数 + 1？
        if res < 1.0:
            res = res + 60.0
        return res 
    
    def detail(self):
        """详情信息，时间戳显示为可读性好的格式
            %Y-%m-%d %H:%M:%S.%f 显示到毫秒
        """
        return f"{self.name} pid: {self.pid} at {ms2str(self.stamp)} cost: {self.cost}ms totalCost: {self.totalCost}"

class TaskInfo:
    """
        描述一个任务的打点信息
            id： 任务的唯一标识
            name: 任务名字，附加信息
            stages: 任务所有阶段打点信息
            startStage: TaskStage 任务初始阶段
    """
    def __init__(self, id, name, startStage = None):
        self.id = id 
        self.name = name
        self.stages = {}
        if startStage:
            self.stages[startStage.name] = startStage

    def addStages(self, stages):
        """添加一个任务的多个阶段
        """
        for stage in stages:
            self.add(stage)
    
    def add(self, stage):
        """添加任务的一个阶段

        Args:
            stage: TaskStage 类型
        """
        if self.__checkValid(stage):
            if isDetail:
                print(f"add {stage.detail()} into {self.id}")
            self.stages[stage.name] = stage
        else:
            print(f"add {stage.detail()} into {self.name} fail!")
            exit()

    def getStart(self):
        """获取任务开始时间点
        """
        return self.stages["Start"].getStart()

    def getEnd(self):
        """获取任务结束时间点
        """
        return self.stages["End"].getStart() + self.stages["End"].cost / 1000
    
    def computeCost(self):
        """计算每个Stage的耗时情况
        """

        items = self.__sortedStages()
        
        for name, stage in items:
            stage.totalCost = stage.stamp - items[0][1].stamp
            if "End" in name:
                stage.cost = stage.stamp - self.stages[getStartName(name)].stamp

    def summary(self):
        """打印任务的总结信息

        name_id:
            Start at 时间（可读性时间戳）
            End at 时间 cost: xxxms
        """
        print(f"{self.name}_{self.id}")
        
        for _, stage in self.__sortedStages():
            print(f"    {stage.detail()}")

    def gantt(self):
        """生成一个任务的gantt描述

        模板: 
            mpJdCommonRequest_30.477__107ms :done, 20:28:30.548, 107ms
            mpJdCommonRequest_30.477__1825ms :crit, 20:28:30.655, 1825ms
            mpJdCommonRequest_30.477__2044ms :active, 20:28:30.477, 2044ms
        """
        name_prefix = f"{self.name}"
        lines = []
        stages = self.__sortedStages()
        flag = 0
        for _, stage in stages:
            if stage.isEnd():
                status = "active"
                if not stage.isTaskEnd():
                    status = "done" if flag % 2 == 0 else "crit"
                flag = flag + 1
                lines.append(f"{name_prefix}_{stage.name}_{stage.cost} :{status}, {ms2str(stage.stamp-stage.cost)[-12:]}, {stage.cost}ms")
        if lines:
            return "\n".join(lines) + "\n"
        return ""

    def barh(self, axes, index, height, stageColor):
        """生成一个任务的gantt图

        Args:
            axes: 绘图区域
            index: 任务的索引
            height: 矩形的高度
        """
        stages = self.__sortedStages()
        for _, stage in stages:
            if stage.isEnd():
                color = stageColor[stage.name]
                offset = height # 总任务绘制在下方
                if not stage.isTaskEnd():
                    offset = 0.0
                
                # if color in ["gray", "yellow"]:
                #     continue
                # 绘制一个矩形
                if isDetail:
                    print(f"axes.barh([{index * height * 2.0 + offset}], [{stage.cost / 1000}], left=[{stage.getStart()}], height={height}, label='{stage.getStageName()}', color='{color}')")
                
                
                rect = axes.barh([index * height * 2.0 + offset], [stage.cost / 1000], left=[stage.getStart()], height=height,
                            label=stage.getStageName(), color=color)
                axes.bar_label(rect, label_type='center', color="white")

    def __sortedStages(self):
        """获取一个按照Stage时间戳排序的列表

        Return:
            [(name, TaskStage), (name, TaskStage), ...]
        """
        items = list(self.stages.items())
        items.sort(key = lambda e : e[1].stamp)
        return items

    def __checkValid(self, stage):
        """校验 TaskStage 是否合法

            如果是End打点，要确保Start也存在
        Args:
            stage: TaskStage 类型
        Return:
            bool
        """

        # 检查stage是否存在
        if stage.name in self.stages:
            # import pdb; pdb.set_trace()
            print(f"{stage.name} already exist")
            return False
        # 检查Start是否存在
        if stage.name == "End":
            return "Start" in self.stages
        # 检查对应的Start是否存在
        if "End" in stage.name:
            return getStartName(stage.name) in self.stages
        return True

class TaskStageParser:
    """将一行日志信息解析为一个TaskStage并添加到对应的任务中

    关键信息:
        任务id: (tag: 1721353735521)
        任务名字: getStorage
        任务时间戳：07-19 09:48:55.524
        任务阶段信息: 
    """
    def parse(self, line):
        taskId = self.parseID(line)
        taskName = self.parseName(line)
        stamp = self.praseStamp(line)
        pid = self.parsePID(line)
        stageNames = self.parseStageName(line, taskName)
        stages = []
        for stageName in stageNames:
            # 时间戳做了转换
            stages.append(TaskStage(stageName, str2ms(stamp), pid))
        if isDetail:
            print(f"parse {taskId} {taskName} {stages}")
        return (taskId, taskName, stages)

    def parseID(self, line):
        """根据日志行，解析出任务ID
        """
        return ""
    
    def parseName(self, line):
        """根据日志行，解析出任务名称
        """
        return ""
    
    def parsePID(self, line):
        """根据日志行，解析出任务执行PID
        """
        return ""
    
    def parseStageName(self, line, taskName):
        """根据日志行，解析任务StageName
        根据关键字映射对应的Stage名字

        Args:
            line: 日志行
            taskName: 任务名 （获取对应的stage配置信息）

        Return:
            返回匹配的列表
        """
        return []

    def praseStamp(self, line):
        """根据日志行，解析出任务Stage的时间戳
        格式: "2024-07-23 12:00:10.123"
        """
        return f"0000-00-00 00:00:00.000"


class MyTaskStageParser(TaskStageParser):
    """将一行日志信息解析为一个TaskStage并添加到对应的任务中

    关键信息:
        任务id: (tag: 1721353735521)
        任务名字: getStorage
            建个映射表，统一一下任务名称
        任务时间戳：07-19 09:48:55.524
        任务阶段信息: （为了配对， 可能要重复生成Stage信息）
            {
                "Start": ["Start:"],
                "End": ["Before Callback:"],
                # "TaskEnd": ["End:"],
                "BeforeRequestStart": ["request task execute start time"],
                "BeforeRequestEnd": ["request start time"],
                "RequestStart": ["request start time"],
                "RequestEnd": ["request end time"]
            }
    """

    def __init__(self, stageNameTable, nameTable):
        super().__init__()
        self.stageNameTable = stageNameTable
        self.nameTable = nameTable


    def parseID(self, line):
        """根据日志行，解析出任务ID

        当前： (tag: ID)
        """
        return getRange(line, "(tag: ", ")")
    
    def parsePID(self, line):
        """根据日志行，解析出任务执行PID
            13027-13027
        """
        columns = line.split()
        return columns[2]

    def parseName(self, line):
        """根据日志行，解析出任务名称
        当前实现:
            基本都在log行最后一个单词
            例外：
                request任务中间打点是倒数第二个（最后一个是个tag信息）
        """
        columns = line.split()
        name = columns[-1]
        if name.isdigit():
            name = columns[-2]
        return self.__getName(name)

    def __getName(self, name):
        """任务名称映射规则
            {
                "normal": ["name1", "name2"]
            }
        """
        # 改为从配置文件读取
        for key, names in self.nameTable.items():
            if name in names:
                return key
        return name

    def parseStageName(self, line, taskName):
        """根据日志行，解析StageName
        根据关键字映射对应的Stage名字
        当前配置映射关系
        stageNames = {
            "Start": ["Start:"],
            "End": ["Before Callback:"],
            "BeforeRequestStart": ["request task execute start time"],
            "BeforeRequestEnd": ["request start time"],
            "RequestStart": ["request start time"],
            "RequestEnd": ["request end time"]
        }

        Return:
            返回匹配的列表
        """
        matchStages = []
        if taskName not in self.stageNameTable:
            return matchStages
        for name, keys in self.stageNameTable[taskName].items():
            for key in keys:
                if key in line:
                    matchStages.append(name)
        if isDetail:
            print(f"match stage names: {matchStages}")
        return matchStages

    def praseStamp(self, line):
        """根据日志行，解析出任务ID
        """
        columns = line.split()
        return f"2024-{columns[0]} {columns[1]}"

    
class LogParser:
    """
        解析一个日志文件对应的任务打点信息

        logPath: 日志文件路径
        configPath: json配置文件路径
        tasks: 保存解析结果
        taskStageParser: 解析日志行的工具
    """
    def __init__(self, logPath, configPath):
        self.logPath = logPath
        self.tasks = {}
        
        with open(configPath, "r+") as jsonFile:
            self.config = json.load(jsonFile)
        self.taskStageParser = MyTaskStageParser(self.config["stage"], self.config["name"])
    
    def parse(self):
        """解析日志文件，提取任务相关打点信息
        """
        with open(self.logPath, "r+") as file:
            for line in file.readlines():
                id, name, stages = self.taskStageParser.parse(line)

                if len(stages) == 0: # 如果没有解析到跳过，避免崩溃
                    continue
                elif len(stages) == 1 and stages[0].name == "Start":
                    # Create TaskInfo
                    self.tasks[id] = TaskInfo(id, name, startStage=stages[0])
                else:
                    self.tasks[id].addStages(stages)
    
    def compute(self):
        """计算所有任务：耗时 和 总耗时
        """
        for _, task in self.tasks.items():
            task.computeCost()

    def traverse(self, cb, filter=None):
        """遍历任务列表，执行cb函数

        Args:
            cb: (task)->void 回调函数
            filter: 过滤关注的任务 比如 ["mpJdCommonRequest"]
        """
        if filter:
            items = [item for item in self.tasks.items() if item[1].name in filter]
        else:
            items = list(self.tasks.items())
        # 按照任务开始时间排序
        items.sort(key = lambda e : e[1].stages["Start"].stamp)
        for index, (_, task) in enumerate(items):
            cb(task, index)

    def get(self, key):
        """获取配置信息
        
        Args:
            key: 链式键值比如 figure.height
        """
        keys = key.split(".")
        result = self.config
        for tk in keys:
            result = result[tk]
        return result

    def savePicture(self, outPath, taskList):
        """生成图片

        Args:
            outPath: 输出路径
            taskList: 关注任务的列表(任务名)
        """
        fig, ax = plt.subplots(figsize=(self.get("figure.width"), self.get("figure.height")))
        ax.invert_yaxis()
        ax.grid(axis="x", color="gray", linewidth=0.5, linestyle="--")

        filterTasks = [task for _, task in self.tasks.items() if task.name in taskList]
        for index, task in enumerate(filterTasks):
            task.barh(ax, index, self.get("figure.lineHeight"), self.get("stageColor"))

        # 整体任务的总结信息： 总任务个数，总任务开始时间，总任务结束时间
        tasksEnd = [task.getEnd() for task in filterTasks]
        taskSize, taskStart, taskEnd = len(filterTasks), filterTasks[0].getStart(), max(tasksEnd)
        ax.set_yticks([self.get("figure.lineHeight") * 2.0 * i for i in range(0, taskSize)])
        # 任务id + 索引
        ax.set_yticklabels([f"{task.name}_{task.id}_{i}"for i, task in enumerate(filterTasks)])
        # 配置x轴刻度和显示label
        width = taskEnd - taskStart
        xticksCount = self.get("figure.xticksCount")
        ax.set_xticks([taskStart + i * width / xticksCount for i in range(xticksCount)])
        ax.set_xticklabels([f"{(taskStart + i * width / xticksCount):.2f}" for i in range(xticksCount)])

        plt.savefig(outPath, format=self.get("figure.format"), dpi=self.get("figure.dpi"))


def getArgs():
    parser = argparse.ArgumentParser()
    parser.add_argument('-p', '--paths', nargs="*", type=str, help='path to the log file')
    parser.add_argument('-c', '--config', nargs=1, type=str, default='config/config.json', help='path to the config file')
    args = parser.parse_args()
    return args.paths, args.config[0]


def testGetArgs():
    logPaths, configPath = getArgs()
    print(f"paths: {logPaths}")
    print(f"config: {configPath}")

def testLog():
    logPaths, configPath = getArgs()

    taskNames = [
        "TaskName"
    ]

    for logPath in logPaths:
        parser = LogParser(logPath, configPath)
        # 解析日志文件
        parser.parse()
        # 计算每个任务的耗时
        parser.compute()

        # 展示解析任务信息
        parser.traverse(lambda task, index: task.summary(), filter=taskNames)
        # 生成 gantt 格式字符串
        # parser.traverse(lambda task, index: print(f"{task.gantt()}", end=""), filter=taskNames)
        # 直接生成图片
        
        outdir = f"output/{os.path.dirname(logPath)}"
        os.makedirs(outdir, exist_ok=True)
        parser.savePicture(f"{outdir}/{os.path.basename(logPath)}_out.svg", taskList=taskNames)
        
if __name__ == "__main__":
    testLog()
    # testGetArgs()
